"""
@Project : KonwledgeDistilling
@Author  : wxj233
@Time    : 2025/11/6 18:47
@Desc    : 
"""
from torchvision import datasets, transforms
import os
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm


os.environ['TORCHVISION_DATA_MIRROR'] = 'https://mirrors.bfsu.edu.cn/torchvision/'
transform = transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])

testset = datasets.MNIST('./', train=False, download=True, transform=transform)
testData = DataLoader(testset, batch_size=16, shuffle=False)


# model = torch.load("model_teacher.pth")
model = torch.load("model_student.pth")
model.eval()

CUDA = torch.cuda.is_available()
t_batch = tqdm(testData, total=len(testData), unit="batches", leave=False)  # 设置leave=False虽然不会换行，但是会导致进度条跑完后被清除
t_batch.set_description(f"测试")

correct = 0
total = 0
for inputs, labels in t_batch:
    if CUDA:
        inputs, labels = inputs.cuda(), labels.cuda()

    outs = model(inputs)
    _, predicted = torch.max(outs, 1)
    correct += (predicted == labels).sum()
    total += labels.shape[0]

print(f"准确率:{correct/total:.5f}")